{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🔗 Link to Github - https://github.com/facebookresearch/llama/blob/6c7fe276574e78057f917549435a2554000a876d/llama/model.py#L64\n",
    "\n",
    "def apply_rotary_emb(\n",
    "    xq: torch.Tensor,\n",
    "    xk: torch.Tensor,\n",
    "    freqs_cis: torch.Tensor,\n",
    ") -> Tuple[torch.Tensor, torch.Tensor]:\n",
    "    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))\n",
    "    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))\n",
    "    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)\n",
    "    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)\n",
    "    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)\n",
    "    return xq_out.type_as(xq), xk_out.type_as(xk)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "🚀Rotary Embeddings (RopE) is now a part of almost every LLMs\n",
    "\n",
    "Let's go through LlaMA-2 Implementation of RopE - nice and simple from LlaMa's Source code - the `apply_rotary_emb()` method 🔥\n",
    "\n",
    "-----------\n",
    "\n",
    "First, briefly the concept of Rotary Embedding.\n",
    "\n",
    "👉 Rotary Embedding (Rotary Positional Encoding) is a kind of positional encoding, but unlike traditional positional encodings, it tries to encode the relative positions of tokens rather than their absolute positions. This is achieved by applying a rotation operation in the embedding space.\n",
    "\n",
    "RoPE is unique in that it encodes the absolute position m of tokens through a rotation matrix R_(θ, m), and also incorporates the explicit relative position dependency in the self-attention formulation. The idea is to embed the position of a token in a sequence by rotating queries and keys, with a different rotation at each position.\n",
    "\n",
    "The main benefit is, Invariance to Sequence Length: Unlike traditional position embeddings, RoPE does not require a predefined maximum sequence length. It can generate position embeddings on-the-fly for any length of sequences. This makes it much more scalable and adaptable to different tasks.\n",
    "\n",
    "-------\n",
    "\n",
    "📌 Now looking at the method `apply_rotary_emb()` \n",
    "\n",
    "👉 The tensors `xq` and `xk` (representing queries and keys) are reshaped to facilitate the creation of complex numbers. The last dimension of each tensor (representing the embedding dimension) is reshaped into two dimensions, with each pair of values treated as a real and imaginary part of a complex number.\n",
    "\n",
    "👉 `xq.shape[:-1]`: This retrieves all dimensions of `xq` except for the last one. For instance, if `xq` has a shape of `(batch_size, num_heads, sequence_length, head_dim)`, `xq.shape[:-1]` would yield `(batch_size, num_heads, sequence_length)`.\n",
    "\n",
    "👉 `reshape(*xq.shape[:-1], -1, 2)` : : The `reshape` operation changes the shape of `xq` to a new shape defined by the parameters. The `*` operator unpacks the shape tuple, effectively expanding it as separate arguments. The `-1` in reshape is a placeholder that automatically calculates the appropriate size for that dimension based on the total number of elements and the other specified dimensions. The `2` at the end splits the last dimension into two dimensions, which is essential for representing complex numbers (as they consist of real and imaginary parts). \n",
    "\n",
    "For example, if `xq` initially has a shape `(32, 8, 128, 64)` (a common shape in transformer models, with 32 being the batch size, 8 the number of heads, 128 the sequence length, and 64 the dimensionality of each head), after this reshaping, it will have the shape `(32, 8, 128, 32, 2)`. The last dimension of `64` is split into two dimensions: `32` for real parts and `2` for representing them as complex numbers (real and imaginary parts).\n",
    "\n",
    "---------\n",
    "\n",
    "Why `xq` and `xk` are being converted into complex tensors with `view_as_complex` ?\n",
    "\n",
    "📌 Rotary Embeddings involve performing a rotation in the embedding space. This rotation is applied to the queries and keys in the attention mechanism of Transformer models.\n",
    "\n",
    "📌 The rotation in RoPE is fundamentally a complex multiplication. Complex numbers are ideal for representing rotations because multiplying by a complex number corresponds to a rotation in the complex plane.\n",
    "\n",
    "📌 The \"rotation\" in this context is a phase shift in the complex plane for each element in the query and key tensors. This shift varies based on the position in the sequence, thus encoding positional information.\n",
    "\n",
    "📌  The essence of RoPE is that by applying these phase shifts, the dot product (used in the self-attention mechanism) between queries and keys becomes sensitive to their relative positions. It's not merely about measuring \"real angle distance\" but rather about how the phase-shifted dot product correlates with the positional relationships of tokens in the input sequence.\n",
    "\n",
    "📌 By converting `xq` and `xk` into complex tensors, these lines prepare them for such a rotation. The actual rotation is performed later in the code, where these complex tensors (`xq_` and `xk_`) are multiplied by another complex tensor representing the rotation (typically through phase factors, like `freqs_cis` in this code).\n",
    "\n",
    "---------\n",
    "\n",
    "👉 Workings of `view_as_complex`\n",
    "\n",
    "👉 The rotation operation here relies on the application of complex numbers due to their intrinsic rotational properties. In the complex plane, multiplying two complex numbers results in adding their angles, which effectively rotates one complex number by the angle of the other.\n",
    "\n",
    "👉 So the function transforms the last dimension of the input tensors `xq` (query) and `xk` (key) into a complex representation. These tensors initially have real values, with the last dimension corresponding to the dimensions of the vector embeddings.\n",
    "\n",
    "Then with, `xq.float().reshape(*xq.shape[:-1], -1, 2)` reshapes the last dimension of `xq` into two dimensions. Each `(real, imaginary)` pair in the last two dimensions represents one complex number.\n",
    "\n",
    "👉 Then, `torch.view_as_complex` takes these pairs of real numbers and treats them as complex numbers, where the first element of each pair is the real part and the second is the imaginary part."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "-----------\n",
    "\n",
    "👉 `reshape_for_broadcast()` is a function that reshapes the frequency embeddings `freqs_cis` to match the dimensions of the input tensors `xq` and `xk`. This ensures that the rotation operation (which is an element-wise multiplication) can be performed.\n",
    "\n",
    "👉 Here, the rotation operation is applied by performing an element-wise multiplication of the input tensors (`xq` and `xk`) and the frequency embeddings (`freqs_cis`). This operation effectively encodes the relative positional information into the embeddings.\n",
    "\n",
    "👉 Then, the result is reshaped back into real numbers (from complex numbers) and flattened to match the original shape of `xq` and `xk`.\n",
    "\n",
    "--------\n",
    "\n",
    "👉 Finally, the function returns the rotated embeddings in the original data type of the input tensors.\n",
    "\n",
    "🔗 Link to SourceCode in Llama Github - https://github.com/facebookresearch/llama/blob/6c7fe276574e78057f917549435a2554000a876d/llama/model.py#L64"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
